import json
import numpy as np
import faiss
from text2vec import SentenceModel

class QuestionRetrievalEngine:
    def __init__(self, questions_path, vectors_path, index_path):
        # 载入主数据和索引
        with open(questions_path, "r", encoding="utf-8") as f:
            self.questions = json.load(f)
        self.vectors = np.load(vectors_path)
        self.index = faiss.read_index(index_path)
        self.model = SentenceModel("shibing624/text2vec-base-chinese")

    def search(self, tags=None, difficulty=None, exclude_ids=None, query=None, topk=3):
        # 多条件过滤
        filtered = [
            q for q in self.questions
            if (not exclude_ids or q["id"] not in exclude_ids)
            and (not tags or any(tag in q["tags"] for tag in tags))
            and (not difficulty or q["difficulty"] == difficulty)
        ]
        if not filtered:
            return []

        # 过滤后做语义检索
        if query:
            sub_vectors = np.array([q["vector"] for q in filtered], dtype="float32")
            query_vec = self.model.encode([query])
            # 这里用L2距离（越小越相似）
            faiss_index = faiss.IndexFlatL2(sub_vectors.shape[1])
            faiss_index.add(sub_vectors)
            D, I = faiss_index.search(query_vec.astype("float32"), topk)
            results = [filtered[i] for i in I[0]]
            return results
        else:
            # 没有query时，直接返回topk
            return filtered[:topk]
